use crate::ast::*;
use crate::error::{CompilerError, Result};
use crate::utils::mangling;
use inkwell::{builder::Builder, values::PointerValue};
use inkwell::context::Context;
use inkwell::module::Module;
use inkwell::types::{BasicType, BasicTypeEnum};
use inkwell::values::{BasicValueEnum, FunctionValue};
use inkwell::AddressSpace;
use std::collections::{HashMap, HashSet};

mod builtins;
mod control_flow;
mod expressions;
mod maps;
mod memory;
mod patterns;
mod rc;
#[path = "slices_impl/mod.rs"]
mod slices;
mod statements;
mod strings;
mod backend_program;
mod backend_scope;
mod backend_utils;
mod backend_link;

pub use memory::PointerOrigin;

/// LLVM backend for code generation
pub struct LLVMBackend<'ctx> {
    /// LLVM context
    pub context: &'ctx Context,
    /// LLVM module
    pub module: Module<'ctx>,
    /// IR builder
    pub builder: Builder<'ctx>,
    /// Function map for tracking generated functions
    pub function_map: HashMap<String, FunctionValue<'ctx>>,
    /// Variable map for tracking local variables
    pub variable_map: HashMap<String, BasicValueEnum<'ctx>>,
    /// Variable type map for tracking variable types (to distinguish pointer vs value)
    pub variable_type_map: HashMap<String, Type>,
    /// Pointer origin map for memory safety (tracks heap vs stack pointers)
    pub pointer_origin_map: HashMap<String, PointerOrigin>,
    /// Stack slices that need automatic cleanup (variable name -> slice type)
    pub stack_slices: Vec<String>,
    /// Struct declarations map for type lookup
    pub struct_declarations: HashMap<String, StructDecl>,
    /// Enum declarations map for type lookup
    pub enum_declarations: HashMap<String, EnumDecl>,
    /// Trait declarations map for object-safe trait handling
    pub trait_declarations: HashMap<String, TraitDecl>,
    /// Trait implementation map (trait, type) -> impl
    pub trait_impls: HashMap<(String, String), TraitImpl>,
    /// Cached LLVM struct types for trait objects
    pub trait_object_types: HashMap<String, inkwell::types::StructType<'ctx>>,
    /// Cached LLVM struct types for trait vtables
    pub trait_vtable_types: HashMap<String, inkwell::types::StructType<'ctx>>,
    /// Cached LLVM struct types for Rc trait handles (rc_ptr + vtable_ptr)
    pub trait_rc_handle_types: HashMap<String, inkwell::types::StructType<'ctx>>,
    /// Cached LLVM struct types for Weak trait handles (weak_ptr + vtable_ptr)
    pub trait_weak_handle_types: HashMap<String, inkwell::types::StructType<'ctx>>,
    /// Cached LLVM function shims for trait methods
    pub trait_method_shims: HashMap<(String, String, String), FunctionValue<'ctx>>,
    /// Cached vtable globals for (trait, type)
    pub trait_vtable_globals: HashMap<(String, String), PointerValue<'ctx>>,
    /// Function declarations for building closures
    function_decls: HashMap<String, FunctionDecl>,
    /// Current function being generated
    pub current_function: Option<FunctionValue<'ctx>>,
    /// NRC return type of the current function (for try/unwrap semantics)
    current_function_return_type: Option<Type>,
    /// Scope stack for defer statements (each scope has its own defer stack)
    pub scope_stack: Vec<Vec<Expression>>,
    /// Recover handler stack (each scope can have recover handlers, LIFO order)
    /// Each element is a RecoverStmt that can catch panics occurring after its definition
    pub recover_stack: Vec<stmt::RecoverStmt>,
    /// After-recover blocks stack (corresponds to recover_stack)
    /// When panic is caught by recover, execution jumps to the corresponding after_recover block
    /// This allows execution to continue after the recover statement, skipping panic and subsequent code
    pub recover_after_blocks: Vec<inkwell::basic_block::BasicBlock<'ctx>>,
    /// Imported module paths that need to be compiled and linked
    pub imported_module_paths: Vec<String>,
    /// Source file directory for resolving relative imports
    pub source_directory: std::path::PathBuf,
    /// Function names that have overloaded definitions and need mangling
    overloaded_function_names: HashSet<String>,
    /// Cached closure shims for named functions
    function_closure_shims: HashMap<String, FunctionValue<'ctx>>,
}

impl<'ctx> LLVMBackend<'ctx> {
    fn function_symbol_name(&self, func: &FunctionDecl) -> String {
        if func.body.is_some() && self.overloaded_function_names.contains(&func.name) {
            mangling::mangle_function_name(&FunctionSignature::from_decl(func))
        } else {
            func.name.clone()
        }
    }

    /// Create a new LLVM backend
    pub fn new(context: &'ctx Context, module_name: &str) -> Result<Self> {
        let module = context.create_module(module_name);
        let builder = context.create_builder();

        // 不再需要JIT执行引擎，只生成可执行文件

        Ok(Self {
            context,
            module,
            builder,
            function_map: HashMap::new(),
            variable_map: HashMap::new(),
            variable_type_map: HashMap::new(),
            pointer_origin_map: HashMap::new(),
            stack_slices: Vec::new(),
            struct_declarations: HashMap::new(),
            enum_declarations: HashMap::new(),
            trait_declarations: HashMap::new(),
            trait_impls: HashMap::new(),
            trait_object_types: HashMap::new(),
            trait_vtable_types: HashMap::new(),
            trait_rc_handle_types: HashMap::new(),
            trait_weak_handle_types: HashMap::new(),
            trait_method_shims: HashMap::new(),
            trait_vtable_globals: HashMap::new(),
            function_decls: HashMap::new(),
            current_function: None,
            current_function_return_type: None,
            scope_stack: Vec::new(),
            recover_stack: Vec::new(),
            recover_after_blocks: Vec::new(),
            imported_module_paths: Vec::new(),
            source_directory: std::env::current_dir()
                .unwrap_or_else(|_| std::path::PathBuf::from(".")),
            overloaded_function_names: HashSet::new(),
            function_closure_shims: HashMap::new(),
        })
    }

    /// Set the source file directory for resolving relative imports
    pub fn set_source_directory(&mut self, dir: std::path::PathBuf) {
        self.source_directory = dir;
    }

    /// Get the LLVM type for an enum (tagged union)
    /// Returns struct { i32 tag, [N x i8] data } where N is size of largest variant
    fn get_enum_llvm_type(&self, enum_name: &str) -> Result<BasicTypeEnum<'ctx>> {
        let enum_decl = self
            .enum_declarations
            .get(enum_name)
            .ok_or_else(|| CompilerError::internal(&format!("Enum '{}' not found", enum_name)))?;

        // Check if this is a simple enum (all unit variants)
        let has_data_variants = enum_decl.variants.iter().any(|v| v.data.is_some());

        if !has_data_variants {
            // Simple enum: just use i32 for the tag
            return Ok(self.context.i32_type().into());
        }

        // Tagged union: calculate the maximum size needed for data
        let mut max_size: usize = 0;

        for variant in &enum_decl.variants {
            if let Some(data) = &variant.data {
                let variant_size = match data {
                    EnumVariantData::Tuple(types) => {
                        // Calculate total size of all fields in tuple
                        let mut total_size = 0;
                        for ty in types {
                            let llvm_type = self.nrc_type_to_llvm_type(&Some(ty.clone()))?;
                            total_size += self.get_llvm_type_size(llvm_type);
                        }
                        total_size
                    }
                    EnumVariantData::Struct(fields) => {
                        // Calculate total size of all fields in struct
                        let mut total_size = 0;
                        for field in fields {
                            let llvm_type =
                                self.nrc_type_to_llvm_type(&Some(field.field_type.clone()))?;
                            total_size += self.get_llvm_type_size(llvm_type);
                        }
                        total_size
                    }
                };
                max_size = max_size.max(variant_size);
            }
        }

        // Create tagged union struct: { i32 tag, [max_size x i8] data }
        let tag_type = self.context.i32_type();
        let data_type = self.context.i8_type().array_type(max_size as u32);

        // Use non-packed struct to ensure proper alignment
        let tagged_union = self.context.struct_type(
            &[tag_type.into(), data_type.into()],
            false, // packed = false
        );

        Ok(tagged_union.into())
    }

    /// Get LLVM type for Result<T, E>
    /// Returns tagged union: struct { i32 tag, [max(sizeof(T), sizeof(E)) x i8] data }
    /// tag = 0 for Ok, tag = 1 for Err
    fn get_result_llvm_type(
        &self,
        result_type: &crate::ast::types::ResultType,
    ) -> Result<BasicTypeEnum<'ctx>> {
        // Get LLVM types for Ok and Err types
        let ok_type = self.nrc_type_to_llvm_type(&Some(*result_type.ok_type.clone()))?;
        let err_type = self.nrc_type_to_llvm_type(&Some(*result_type.err_type.clone()))?;

        // Calculate sizes
        let ok_size = self.get_llvm_type_size(ok_type);
        let err_size = self.get_llvm_type_size(err_type);

        // Use the maximum size for the data array
        let max_size = ok_size.max(err_size);

        // Create tagged union struct: { i32 tag, [max_size x i8] data }
        let tag_type = self.context.i32_type();
        let data_type = self.context.i8_type().array_type(max_size as u32);

        // Use non-packed struct to ensure proper alignment
        let tagged_union = self.context.struct_type(
            &[tag_type.into(), data_type.into()],
            false, // packed = false
        );

        Ok(tagged_union.into())
    }

    fn get_option_llvm_type(
        &self,
        option_type: &crate::ast::types::OptionType,
    ) -> Result<BasicTypeEnum<'ctx>> {
        // Get LLVM type for Some type
        let some_type = self.nrc_type_to_llvm_type(&Some(*option_type.some_type.clone()))?;

        // Calculate size
        let some_size = self.get_llvm_type_size(some_type);

        // Data field is an array of bytes large enough to hold the value
        let data_size = some_size.max(1); // At least 1 byte

        // Create tagged union struct: { i32 tag, [data_size x i8] data }
        let tag_type = self.context.i32_type();
        let data_type = self.context.i8_type().array_type(data_size as u32);

        // Use non-packed struct to ensure proper alignment
        let tagged_union = self.context.struct_type(
            &[tag_type.into(), data_type.into()],
            false, // packed = false
        );

        Ok(tagged_union.into())
    }

    /// Get the size of an LLVM basic type in bytes (for enum layout calculation)
    fn get_llvm_type_size(&self, llvm_type: BasicTypeEnum<'ctx>) -> usize {
        match llvm_type {
            BasicTypeEnum::IntType(int_type) => {
                // For i1 (bool), we need at least 1 byte
                let bits = int_type.get_bit_width();
                ((bits + 7) / 8) as usize // Round up to nearest byte
            }
            BasicTypeEnum::FloatType(float_type) => {
                match float_type.get_context().f32_type() == float_type {
                    true => 4,
                    false => 8, // f64
                }
            }
            BasicTypeEnum::PointerType(_) => 8, // 64-bit pointer
            BasicTypeEnum::ArrayType(array_type) => {
                let element_size = self.get_llvm_type_size(array_type.get_element_type());
                element_size * array_type.len() as usize
            }
            BasicTypeEnum::StructType(struct_type) => {
                let mut total = 0;
                for field_type in struct_type.get_field_types() {
                    total += self.get_llvm_type_size(field_type);
                }
                total
            }
            _ => 8, // default
        }
    }

    /// Extract tag field from enum value (for comparison and pattern matching)
    pub(super) fn extract_enum_tag(
        &mut self,
        enum_value: BasicValueEnum<'ctx>,
        enum_type_name: &str,
    ) -> Result<BasicValueEnum<'ctx>> {
        // Check if this is a simple enum (no data) or tagged union
        let enum_decl = self.enum_declarations.get(enum_type_name).ok_or_else(|| {
            CompilerError::internal(&format!("Enum '{}' not found", enum_type_name))
        })?;

        let has_data_variants = enum_decl.variants.iter().any(|v| v.data.is_some());

        if !has_data_variants {
            // Simple enum: the value IS the tag
            return Ok(enum_value);
        }

        // Tagged union: extract the tag field (index 0)
        if let BasicValueEnum::StructValue(struct_val) = enum_value {
            // Extract first field (tag)
            let tag_value = self.builder.build_extract_value(struct_val, 0, "tag")?;
            Ok(tag_value.into())
        } else {
            Err(CompilerError::internal(
                "Expected struct value for enum with data",
            ))
        }
    }

    /// Extract data from enum variant (for pattern matching)
    /// Returns a pointer to the data array
    pub(super) fn extract_enum_data_ptr(
        &mut self,
        enum_ptr: inkwell::values::PointerValue<'ctx>,
        enum_type_name: &str,
    ) -> Result<inkwell::values::PointerValue<'ctx>> {
        let enum_llvm_type = self.get_enum_llvm_type(enum_type_name)?;

        // Get pointer to data field (index 1)
        let data_ptr = unsafe {
            self.builder.build_gep(
                enum_llvm_type,
                enum_ptr,
                &[
                    self.context.i32_type().const_zero(),
                    self.context.i32_type().const_int(1, false),
                ],
                "data_ptr",
            )?
        };

        Ok(data_ptr)
    }

    /// Generate enum variant constructor (e.g., Result::Ok(42))
    /// Creates a tagged union value with the discriminant and data
    pub(super) fn generate_enum_variant_constructor(
        &mut self,
        enum_name: &str,
        variant_name: &str,
        arguments: &[Expression],
    ) -> Result<BasicValueEnum<'ctx>> {
        let enum_decl = self
            .enum_declarations
            .get(enum_name)
            .ok_or_else(|| CompilerError::internal(&format!("Enum '{}' not found", enum_name)))?
            .clone();

        // Find the variant and its index
        let (variant_index, variant) = enum_decl
            .variants
            .iter()
            .enumerate()
            .find(|(_, v)| v.name == variant_name)
            .ok_or_else(|| {
                CompilerError::internal(&format!(
                    "Variant '{}::{}' not found",
                    enum_name, variant_name
                ))
            })?;

        // Get the enum LLVM type
        let enum_llvm_type = self.get_enum_llvm_type(enum_name)?;

        // Check if this is a simple enum (no data)
        let has_data_variants = enum_decl.variants.iter().any(|v| v.data.is_some());

        if !has_data_variants {
            // Simple enum: just return the tag value
            let tag_value = self
                .context
                .i32_type()
                .const_int(variant_index as u64, false);
            return Ok(tag_value.into());
        }

        // Tagged union: create the struct value
        // Allocate space for the enum
        let enum_alloca = self.builder.build_alloca(enum_llvm_type, "enum_temp")?;

        // Set the tag field (index 0)
        let tag_value = self
            .context
            .i32_type()
            .const_int(variant_index as u64, false);
        let tag_ptr = unsafe {
            self.builder.build_gep(
                enum_llvm_type,
                enum_alloca,
                &[
                    self.context.i32_type().const_zero(),
                    self.context.i32_type().const_zero(),
                ],
                "tag_ptr",
            )?
        };
        self.builder.build_store(tag_ptr, tag_value)?;

        // Set the data field (index 1) if this variant has data
        if let Some(variant_data) = &variant.data {
            let data_ptr = unsafe {
                self.builder.build_gep(
                    enum_llvm_type,
                    enum_alloca,
                    &[
                        self.context.i32_type().const_zero(),
                        self.context.i32_type().const_int(1, false),
                    ],
                    "data_ptr",
                )?
            };

            match variant_data {
                EnumVariantData::Tuple(types) => {
                    // Verify argument count
                    if arguments.len() != types.len() {
                        return Err(CompilerError::internal(&format!(
                            "Variant '{}::{}' expects {} arguments, got {}",
                            enum_name,
                            variant_name,
                            types.len(),
                            arguments.len()
                        )));
                    }

                    // Store each argument into the data array
                    let mut offset = 0;
                    for (i, (ty, arg)) in types.iter().zip(arguments.iter()).enumerate() {
                        let arg_value = self.generate_expression(arg)?;
                        let arg_llvm_type = self.nrc_type_to_llvm_type(&Some(ty.clone()))?;

                        // Calculate pointer to this field in the data array
                        let field_ptr = unsafe {
                            self.builder.build_gep(
                                self.context.i8_type(),
                                data_ptr,
                                &[self.context.i32_type().const_int(offset as u64, false)],
                                &format!("field_{}_ptr", i),
                            )?
                        };

                        // Cast to the correct type and store
                        let typed_ptr = self.builder.build_pointer_cast(
                            field_ptr,
                            arg_llvm_type.ptr_type(AddressSpace::default()),
                            &format!("field_{}_typed_ptr", i),
                        )?;

                        self.builder.build_store(typed_ptr, arg_value)?;

                        offset += self.get_llvm_type_size(arg_llvm_type);
                    }
                }
                EnumVariantData::Struct(_fields) => {
                    return Err(CompilerError::internal(
                        "Struct-style enum variants not yet implemented",
                    ));
                }
            }
        } else if !arguments.is_empty() {
            return Err(CompilerError::internal(&format!(
                "Variant '{}::{}' is a unit variant and does not accept arguments",
                enum_name, variant_name
            )));
        }

        // Load the complete enum value
        let enum_value = self
            .builder
            .build_load(enum_llvm_type, enum_alloca, "enum_value")?;
        Ok(enum_value)
    }

    /// Generate global constants for enum variants
    /// For unit variants (no data), create global constants
    /// For data variants, the tag value will be used at construction time
    fn generate_enum_constants(&mut self, enum_decl: &EnumDecl) -> Result<()> {
        for (index, variant) in enum_decl.variants.iter().enumerate() {
            // Create global constant for the tag value
            // Name format: EnumName_VariantName (e.g., Status_Pending, Result_Ok)
            let const_name = format!("{}_{}", enum_decl.name, variant.name);
            let const_value = self.context.i32_type().const_int(index as u64, false);

            // Add global constant to module
            let global = self
                .module
                .add_global(self.context.i32_type(), None, &const_name);
            global.set_initializer(&const_value);
            global.set_constant(true);

            // For unit variants (no data), also add to variable map so they can be used directly
            if variant.data.is_none() {
                self.variable_map
                    .insert(variant.name.clone(), const_value.into());
            }
            // For data variants, we don't add to variable_map because they need constructor calls
        }

        Ok(())
    }

}
